import cv2
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
# yoloV3的预测器
from core.predicter import Predictor
import tkinter as tk
from tkinter import filedialog, ttk
from PIL import Image, ImageTk, ImageDraw, ImageFont
import numpy as np
import os
# 完全禁用GPU以避免CUDA相关错误
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

# coco数据集中的类别信息
classes = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
           'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
           'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
           'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
           'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
           'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
           'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork',
           'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli',
           'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant',
           'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
           'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book',
           'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']


def visualize_detections(image, boundings):
    """
    将检测结果可视化到图像上。
    :param image: PIL Image 对象
    :param boundings: 检测结果列表
    :return: 可视化后的 PIL Image 对象
    """
    draw = ImageDraw.Draw(image)
    try:
        font = ImageFont.truetype("arial.ttf", 12)
    except IOError:
        font = ImageFont.load_default()

    for bounding in boundings:
        if bounding[4].numpy() > 0.5:  # 过滤低置信度检测结果
            x1, y1, x2, y2 = bounding[0].numpy(), bounding[1].numpy(), bounding[2].numpy(), bounding[3].numpy()
            label_id = int(bounding[5].numpy())
            label = classes[label_id]
            score = bounding[4].numpy()

            draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
            text = f"{label}: {score:.2f}"
            draw.text((x1, y1 - 15), text, fill="red", font=font)
    return image


class ImageViewer:
    def __init__(self, root):
        self.root = root
        self.root.title("图像查看器")
        self.root.geometry("1600x800")  # 调整窗口大小以容纳两个图像框

        # 创建按钮
        self.btn_select = tk.Button(root, text="选择图像", command=self.select_image)
        self.btn_select.pack(pady=10)

        # 创建滚动条
        self.progress_bar = ttk.Progressbar(root, orient="horizontal", length=300, mode="determinate")
        self.progress_bar.pack(pady=10)

        # 创建原始图像显示标签
        self.original_img_label = tk.Label(root)
        self.original_img_label.pack(side=tk.LEFT, padx=10)

        # 创建检测结果显示标签
        self.result_img_label = tk.Label(root)
        self.result_img_label.pack(side=tk.RIGHT, padx=10)

        # 创建用于显示检测目标分类的文本框
        self.detected_objects_text = tk.Text(root, height=20, width=30)
        self.detected_objects_text.pack(pady=10)

    def select_image(self):
        # 清空上一次检测结果的显示
        self.detected_objects_text.delete(1.0, tk.END)

        # 打开文件选择对话框
        file_path = filedialog.askopenfilename(
            title="选择图像文件",
            filetypes=[("图像文件", "*.png;*.jpg;*.jpeg;*.bmp")]
        )

        if not file_path:  # 用户取消选择
            return

        # 更新滚动条为 20% 进度
        self.progress_bar['value'] = 20
        self.root.update_idletasks()

        # 加载原始图像
        original_img = Image.open(file_path)
        original_img.thumbnail((780, 580))  # 限制最大显示尺寸（保留宽高比）
        self.original_tk_img = ImageTk.PhotoImage(original_img)
        self.original_img_label.config(image=self.original_tk_img)

        # 更新滚动条为 40% 进度
        self.progress_bar['value'] = 40
        self.root.update_idletasks()

        # Convert PIL image to numpy array
        img_array = np.array(original_img)

        # 更新滚动条为 60% 进度
        self.progress_bar['value'] = 60
        self.root.update_idletasks()

        # 进行目标检测
        predictor = Predictor(class_num=80, yolov3="weights/yolov3.h5")
        try:
            boundings = predictor.predict(img_array)
        except Exception as e:
            print(f"目标检测出错: {e}")
            return

        # 更新滚动条为 80% 进度
        self.progress_bar['value'] = 80
        self.root.update_idletasks()

        # 可视化检测结果
        result_img = visualize_detections(original_img.copy(), boundings)
        result_img.thumbnail((780, 580))  # 限制最大显示尺寸（保留宽高比）
        self.result_tk_img = ImageTk.PhotoImage(result_img)
        self.result_img_label.config(image=self.result_tk_img)

        # 统计检测到的目标分类
        detected_objects = {}
        for bounding in boundings:
            if bounding[4].numpy() > 0.5:  # 过滤低置信度检测结果
                label_id = int(bounding[5].numpy())
                label = classes[label_id]
                if label in detected_objects:
                    detected_objects[label] += 1
                else:
                    detected_objects[label] = 1

        # 在文本框中显示检测到的目标分类
        for obj, count in detected_objects.items():
            self.detected_objects_text.insert(tk.END, f"{obj}: {count}\n")

        # 创建一个新窗口显示所有检测到的目标
        all_objects_window = tk.Toplevel(self.root)
        all_objects_window.title("所有检测到的目标")

        # 创建一个画布用于显示多个目标图像
        canvas = tk.Canvas(all_objects_window)
        canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        # 创建一个滚动条
        scrollbar = tk.Scrollbar(all_objects_window, orient=tk.VERTICAL, command=canvas.yview)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)

        canvas.configure(yscrollcommand=scrollbar.set)
        canvas.bind('<Configure>', lambda e: canvas.configure(scrollregion=canvas.bbox("all")))

        # 创建一个框架用于放置目标图像
        frame = tk.Frame(canvas)
        canvas.create_window((0, 0), window=frame, anchor="nw")

        row = 0
        col = 0
        max_cols = 3  # 每行最多显示 3 个目标

        for bounding in boundings:
            if bounding[4].numpy() > 0.5:
                x1, y1, x2, y2 = int(bounding[0].numpy()), int(bounding[1].numpy()), int(bounding[2].numpy()), int(bounding[3].numpy())
                label_id = int(bounding[5].numpy())
                label = classes[label_id]

                # 裁剪出目标图像
                cropped_img = original_img.crop((x1, y1, x2, y2))
                cropped_img.thumbnail((200, 200))  # 调整目标图像大小

                # 显示裁剪后的图像
                cropped_tk_img = ImageTk.PhotoImage(cropped_img)
                img_label = tk.Label(frame, image=cropped_tk_img)
                img_label.image = cropped_tk_img
                img_label.grid(row=row, column=col, padx=5, pady=5)

                # 显示目标类别标签
                label_text = tk.Label(frame, text=label)
                label_text.grid(row=row + 1, column=col, padx=5, pady=5)

                col += 1
                if col >= max_cols:
                    col = 0
                    row += 2

        # 更新滚动条为 100% 进度
        self.progress_bar['value'] = 100
        self.root.update_idletasks()


if __name__ == "__main__":
    root = tk.Tk()
    app = ImageViewer(root)
    root.mainloop()
